import torch
from tqdm import tqdm
from DEAD.AutoDecoder.Generator.GDS import load_npz
import numpy as np

class GrainDatasetOnGPU():
    def __init__(self,data,batch_size):
        # 将所有数据转移到 GPU 上
        self.list_index = data[0].cuda()
        self.list_burning_surface_coordinate = data[1].cuda()
        self.list_cavity_region_coordinate = data[2].cuda()
        self.list_cavity_region_sdf = data[3].cuda()
        self.list_solid_region_coordinate = data[4].cuda()
        self.list_solid_region_sdf = data[5].cuda()
        self.list_lz = data[6].cuda()
        self.list_n_slots = data[7].cuda()
        self.list_m = data[8].cuda()

        # 计算数据集大小和批次数
        self.length = len(self.list_index)
        self.batch_size=batch_size
        self.num_batches = (self.length + batch_size - 1) // batch_size # 向上取整 
    
    def get_indices_random(self):
        # 需要打乱数据
        return torch.randperm(self.length).cuda()

    def get_indices(self):
        # 不打乱顺序
        return torch.arange(self.length).cuda()
    
    def get_batch(self, batch_index, indices):
        # 生成批次索引
        start = batch_index * self.batch_size
        end = min(start + self.batch_size, self.length)  # 确保不超出数据集范围
        batch_indices = indices[start:end]
        # 自定义批处理逻辑
        return (
            self.list_index[batch_indices],
            self.list_burning_surface_coordinate[batch_indices],
            self.list_cavity_region_coordinate[batch_indices],
            self.list_cavity_region_sdf[batch_indices],
            self.list_solid_region_coordinate[batch_indices],
            self.list_solid_region_sdf[batch_indices]
        )

    def get_geometry_vector(self, index):
        return torch.stack((self.list_lz[index], self.list_n_slots[index], self.list_m[index]), dim=0).to(torch.float32)

def load_data_tensor(data_path: str, data_num:int, data_index_start:int=0):  # 载入数据
    list_index = []
    list_burning_surface_coordinate = []
    list_cavity_region_coordinate = []
    list_cavity_region_sdf = []
    list_solid_region_coordinate = []
    list_solid_region_sdf = []
    list_lz = []
    list_n_slots = []
    list_m = []

    for i in tqdm(range(data_index_start, data_index_start + data_num), desc='loading data', colour='green', dynamic_ncols=True):
        npz_path = f'{data_path}/{i}'
        gds = load_npz(npz_path=npz_path, need_compress=True)
        list_index.append(i-data_index_start)
        list_burning_surface_coordinate.append(gds.burning_surface_coordinate)
        list_cavity_region_coordinate.append(gds.cavity_region_coordinate)
        list_cavity_region_sdf.append(gds.cavity_region_sdf)
        list_solid_region_coordinate.append(gds.solid_region_coordinate)
        list_solid_region_sdf.append(gds.solid_region_sdf)
        list_lz.append(gds.lz)
        list_n_slots.append(gds.n_slots)
        list_m.append(gds.m)

    # 将numpy列表转换为PyTorch张量
    list_index = torch.tensor(np.stack(list_index, axis=0))
    list_burning_surface_coordinate = torch.tensor(
        np.stack(list_burning_surface_coordinate, axis=0), dtype=torch.float32)
    list_cavity_region_coordinate = torch.tensor(
        np.stack(list_cavity_region_coordinate, axis=0), dtype=torch.float32)
    list_cavity_region_sdf = torch.tensor(
        np.stack(list_cavity_region_sdf, axis=0), dtype=torch.float32)
    list_solid_region_coordinate = torch.tensor(
        np.stack(list_solid_region_coordinate, axis=0), dtype=torch.float32)
    list_solid_region_sdf = torch.tensor(
        np.stack(list_solid_region_sdf, axis=0), dtype=torch.float32)
    list_lz = torch.tensor(np.stack(list_lz, axis=0), dtype=torch.float32)
    list_n_slots = torch.tensor(np.stack(list_n_slots, axis=0))
    list_m = torch.tensor(np.stack(list_m, axis=0), dtype=torch.float32)

    return list_index, \
        list_burning_surface_coordinate, \
        list_cavity_region_coordinate, \
        list_cavity_region_sdf, \
        list_solid_region_coordinate, \
        list_solid_region_sdf, \
        list_lz, \
        list_n_slots,\
        list_m